import torch
import torch.nn as nn
import torch.nn.functional as F

"""
class FocalLoss(nn.Module):

    def __init__(self, gamma=2, eps=1e-7):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.eps = eps
        self.ce = nn.CrossEntropyLoss()
        self.softmax = F.softmax

    def forward(self, input, target):
        input = self.softmax(input)
        logp = self.ce(input, target.long())
        p = torch.exp(-logp)
        loss = (1 - p) ** self.gamma * logp
        return loss.mean()

"""
class FocalLoss():
    def __init__(self, gamma=2, alpha=0.25):
        self.gamma = gamma
        self.alpha = alpha

    def __call__(self, input, target):
        input = F.softmax(input)
        criterion = nn.CrossEntropyLoss()
        criterion = criterion.cuda()

        logpt = -criterion(input, target.long())
        pt = torch.exp(logpt)
        logpt *= self.alpha
        loss = -((1 - pt) ** self.gamma) * logpt

        loss /= input.size(0)

        return loss
